# Various patran helper objects
import fortranformat as ff
import collections
import re

# Valid warp3d result file rules:
initial = ['w']
location = ['n', 'e']
types = ['r','e','v','t','a','d','s','m']
suffixes = ['_text', '_stream']

mats = []

# Big lookup table: for each description (i.e. e = strains) given a huge
# list of the field names
results_lookup = {
		'r': ['Reactionx', 'Reactiony', 'Reactionz'],
		'e': ['Strain-xx', 'Strain-yy', 'Strain-zz', 'Strain-xy', 'Strain-yz',
			'Strain-xz', 'Effective mises strain', 'Strain invariant 1',
			'Strain invariant 2', 'Strain invariant 3', 'Strain principle (min)',
			'Strain principle (inter)', 'Strain principle (max)',
			'Strain dir cosine l1', 'Strain dir cosine m1', 'Strain dir cosine n1',
			'Strain dir cosine l2', 'Strain dir cosine m2', 'Strain dir cosine n2',
			'Strain dir cosine l3', 'Strain dir cosine m3', 'Strain dir cosine n3'],
		'v': ['Velocityx', 'Velocityy', 'Velocityz'],
		't': ['Temperature', 'Reference temperature'],
		'a': ['Accelerationx', 'Accelerationy', 'Accelerationz'],
		'd': ['Displacementx', 'Displacementy', 'Displacementz'],
		's': ['Stress-xx', 'Stress-yy', 'Stress-zz', 'Stress-xy', 'Stress-yz',
			'Stress-xz', 'Stress work density', 'mises stress', 'C1-Mat val',
			'C2-Mat val', 'C3-Mat val', 'Stress invariant 1', 'Stress invariant 2',
			'Stress invariant 3', 'Stress principle (min)', 'Stress principle (inter)',
			'Stress principle (max)',
			'Stress dir cosine l1', 'Stress dir cosine m1', 'Stress dir cosine n1',
			'Stress dir cosine l2', 'Stress dir cosine m2', 'Stress dir cosine n2',
			'Stress dir cosine l3', 'Stress dir cosine m3', 'Stress dir cosine n3']
		}

# Don't change, just compiles the regular expression
reor = lambda s: '('+'|'.join(s)+')'
restr = ''.join([reor(initial), reor(location), reor(types), r'(\d*)',
	reor(suffixes),reor(mats)+'?',r'(?:.gz)?','$'])
reg = re.compile(restr)

def recompile():
	"""
		Blargh: after changing the list of valid mats you need to recompile the
		stupid regex...
	"""
	reor = lambda s: '('+'|'.join(s)+')'
	restr = ''.join([reor(initial), reor(location), reor(types), r'(\d*)',
		reor(suffixes),reor(mats)+'?',r'(?:.gz)?','$'])
	reg = re.compile(restr)
	return reg

def valid_result_file(fname):
	"""
		Determine if a file fname is a valid Warp3D results file, according to 
		the above rules.
	"""
	if reg.search(fname):
		return True
	else:
		return False

def parse_result_fname(fname):
	"""
		Separate a result file name into the location, type, step #, and file type
	"""
	mtchs = reg.match(fname)
	if mtchs.group(6):
		typ = mtchs.group(3)+mtchs.group(6)
	else:
		typ = mtchs.group(3)

	if not mtchs:
		raise ValueError("Filename %s is not a valid results file!" % fname)

	return (mtchs.group(2), typ, int(mtchs.group(4)), 
			mtchs.group(5).lstrip('_'))

# Element type definitions
pt = {
		2: 'TRUSS',
		3: 'TRIANGLE',
		4: 'QUAD',
		5: 'TETRA',
		7: 'WEDGE',
		8: 'HEX'}

patran_types = collections.defaultdict(lambda: 'UNKNOWN', pt)

# Fortran packet definitions.  Dictionary linking the packet number to 
# the relevant data formats.  If KC > # of formats just repeat the last one.
packets = { 	25: ('(20A4)',),
							26: ('(3A4,2A4,3A4)',),
							1:  ('(3E16.9)','(I1,1A1,I8,I8,I8,2X,6I1)'),
							2:  ('(I8,I8,I8,I8,3E16.9)', '(10I8)'),
							3:  ('(5E16.9)',),
							4:	('raw',),
							5:	('(5E16.9)',),
							6:  ('(I1,I1,I1,6I1,8I1,I2)', '(5E16.9)'),
							7:  ('(I8,6I1)', '(5E16.9)'),
							8:  ('(I8,6I1)', '(5E16.9)'),
							10: ('(E16.9',),
							11: ('(E16.9',),
							14: ('(3A12)','(2I8,E16.9)','2(2I8,E16.9)'),
							15: ('(E16.9)',),
							16: ('(I1,1X,8I1)','(5E16.9)'),
							17: ('(I1,1X,8I1)','(5E16.9)'),
							18: ('(I1,1X,8I1)','(5E16.9)'),
							19: ('(6I8,2X,8I1)',),
							21: ('(A12)','(10I8)'),
							31: ('(3E16.9)',),
							32: 3*('(5E16.9,5E16.9,2E16.9,2I8)',),
							33: 10*('9(5E16.9/5), 3E16.9/2E16.9,4I8',),
							42: ('(I8,3I8,I8,5I8)','(I8,I8,7X,1A1,7I8)','(3E16.9)','(3E16.9)','(8I8, I8, I8)','(3E16.9, I8, I8)'),
							43: ('(I8,3I8,I8,5I8)','(I8,I8,7X,1A1,7I8)','(3E16.9)','(3E16.9)','(8I8, I8, I8)','(3E16.9, I8, I8)'),
							44: ('(I8,3I8,I8,5I8)','(I8,I8,7X,1A1,7I8)','(3E16.9)','(3E16.9)','(8I8, I8, I8)','(3E16.9, I8, I8)'),
							45: ('(6I8)','(10I8)'),
							99: tuple()
							}

class PatranResultsDesc(object):
	"""
		Descriptor class for a set of results files at a common location and type.
		Contains step numbers, full paths of datafiles, and file types.
	"""
	def __init__(self, loc, rtype):
		self.loc = loc
		self.rtype = rtype
		self.steps = []
		self.fpaths = []
		self.ftypes = []
		self.labels = []
	
	def __len__(self):
		return len(self.labels)

	@property
	def description(self):
		"""
			Return the 2 character string describing the type.
		"""
		return self.loc+self.rtype

	def add_file_at_step(self, step, ftype, path):
		"""
			Add datafile with path at step #.
		"""
		if step in self.steps:
			raise ValueError("Results file for type %s at step %i already exists."
					% (self.description,step))

		# Find the correct place to insert, this gives the index beyond the spot
		# to insert
		i = next((i for i,v in enumerate(self.steps) if v > step), len(self.steps))
		self.steps.insert(i,step)
		self.ftypes.insert(i,ftype)
		self.fpaths.insert(i,path)

class PatranNeutralIt(object):
	"""
		An iterator which reads packets off of a patran neutral file, in order.

		Returns packets in format:
		( (#, ID, IV, KC, N1, N2, N3, N4, N5), [ (each), (data), (card) ])
		Where each data card is parsed by the appropriate card parser.
	"""
	def __init__(self, filename, writeout=None):
		"""
			filename -- patran neutral file
		"""
		self.file = open(filename)
		self.n = 0
		self.p = 0
		self.writeout = writeout

	def __iter__(self):
		return self

	def next(self):
		"""
			Read the next packet.  Raise StopIteration if at the end of the file,
			raise an error if the format of the next packet doesn't match.
		"""
		line = self.file.readline()
		if not line:
			raise StopIteration()
		self.n += 1

		if self.writeout and (self.p % self.writeout) == 0 and self.p != 0:
			print("\tRead %i packets" % self.p)

#		Try to parse as a headerline
		try:
			header_parser = ff.FortranRecordReader('(I2,8I8)')
			header_data = header_parser.read(line)
		except ValueError:
			raise ValueError("Line %i does not appear to be a valid packet header.\n"
					"Line: %s" % (self.n, line))

#		Grab the packet type and the number of data cards
		pn = header_data[0]
		kc = header_data[3]

		if pn not in packets:
			raise ValueError("Unknown packet type %i on line %i." % (pn, self.n))

		formats = packets[pn]

# 	Just stop if you see a stop packet
		if pn == 99:
			return (header_data, [])

		data = []
#		For each data packet...
		for i in range(kc):
			line = self.file.readline()
			if not line:
				raise ValueError("EOF occurred before all data packets read!")
			self.n += 1
#			Assume that the last format takes care of the rest of the packets
			if i >= len(formats):
				f = formats[-1]
			else:
				f = formats[i]
			
			if f == 'raw':
				data.append(line)
			else:
#				Try parsing out the data
				try:
					data_parser = ff.FortranRecordReader(f)
					cpacket_data = data_parser.read(line)
				except ValueError:
					raise ValueError("Line %i does not appear to be a valid data packet.\n"
							"Line: %s" % (self.n, line))
				data.append(cpacket_data)

		self.p += 1

		return (header_data, data)
